#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import os

from yolox.exp import SidaExp as MyExp
from yolox.data.datasets.sida_mini_coco_classes import SIDA_MINI_COCO_CLASSES as CLASSES


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()

        # ---------------- model config ---------------- #
        self.classes = CLASSES
        self.num_classes = len(self.classes)
        self.depth = 0.33
        self.width = 0.50
        self.act = "relu"

        # ---------------- dataloader config ---------------- #
        # set worker to 4 for shorter dataloader init time
        self.data_num_workers = 8
        self.input_size = (640, 640)
        self.random_size = (14, 26)
        self.data_dir = None
        self.train_ann = "train_coco_mini.json"
        self.val_ann = "val_coco.json"
        self.train_img_dir = "train2017"
        self.train_ann_dir = "annotations"
        self.val_img_dir = "val2017"
        self.val_ann_dir = "annotations"

        # --------------- transform config ----------------- #
        self.degrees = 10.0
        self.translate = 0.1
        self.scale = (0.1, 2)
        self.mscale = (0.8, 1.6)
        self.shear = 2.0
        self.perspective = 0.0
        self.enable_mixup = True

        # --------------  training config --------------------- #
        self.warmup_epochs = 2
        self.max_epoch = 24
        self.warmup_lr = 0
        # if basic_lr is None, use basic_lr_per_img*batch_size as basic_lr
        # self.basic_lr = 0.01
        self.basic_lr = None
        self.basic_lr_per_img = 0.01 / 16.0
        self.scheduler = "yoloxwarmcos"
        self.no_aug_epochs = 2
        self.min_lr_ratio = 0.05
        self.ema = True

        self.weight_decay = 5e-4
        self.momentum = 0.9
        self.print_interval = 10
        self.eval_interval = self.max_epoch
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # -----------------  testing config ------------------ #
        self.test_size = (640, 640)
        self.test_conf = 0.01
        self.nmsthre = 0.6

        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
